#include "gtd_common.hpp"
#include "gtd_tracker.hpp"

// WARNING: There is a bug in CUDA 3.0 where textures inside namespaces => "unspecified driver error"
texture<float, 2, cudaReadModeElementType> downSampler;
texture<float, 2, cudaReadModeElementType> gradxSampler;
texture<float, 2, cudaReadModeElementType> gradySampler;
texture<float, 2, cudaReadModeElementType> image1Sampler;
texture<float, 2, cudaReadModeElementType> image2Sampler;
texture<float, 2, cudaReadModeElementType> gradx1Sampler;
texture<float, 2, cudaReadModeElementType> gradx2Sampler;
texture<float, 2, cudaReadModeElementType> grady1Sampler;
texture<float, 2, cudaReadModeElementType> grady2Sampler;

NAMESPACE_BEG(gtd)

#define NORM_THREAD_X 64
#define NORM_BLOCK_X 64
#define NORM_PIXELS_PER_BLOCK (NORM_THREAD_X * NORM_BLOCK_X)

__constant__ uint2 INITIAL_RANGE = { 765, 0 };
__device__ uint2 normalRange;

__global__ void normalizeImage1(const uchar4* src)
{
    __shared__ uint2 blockRange[NORM_BLOCK_X];

    uint base = blockIdx.x * NORM_PIXELS_PER_BLOCK + threadIdx.x;
    
    uint2 range = { 765, 0 };

    #pragma unroll
    for(uint i = 0; i < NORM_THREAD_X; ++i)
    {
        uchar4 pixel = src[base + i * NORM_BLOCK_X];
        uint I = (uint)pixel.x + pixel.y + pixel.z;
        range.x = min(range.x, I);
        range.y = max(range.y, I);
    }
    blockRange[threadIdx.x] = range;

    __syncthreads();

    for(uint i = 2; i <= NORM_BLOCK_X; i <<= 1)
    {
        if((threadIdx.x & (i - 1)) == 0)
        {
            blockRange[threadIdx.x].x = min(blockRange[threadIdx.x].x, 
                blockRange[threadIdx.x + (i >> 1)].x);
            blockRange[threadIdx.x].y = max(blockRange[threadIdx.x].y, 
                blockRange[threadIdx.x + (i >> 1)].y);
        }

        __syncthreads();
    }

    if(threadIdx.x == 0)
    {
        atomicMin(&normalRange.x, blockRange[0].x);
        atomicMax(&normalRange.y, blockRange[0].y);
    }
}

__global__ void normalizeImage2(float* dst, const uchar4* src)
{
    const float scale = 1.0f / (normalRange.y > normalRange.x ? (float)normalRange.y - normalRange.x : 765.0f);
    const float shift = -normalRange.x;

    unsigned base = blockIdx.x * NORM_PIXELS_PER_BLOCK + threadIdx.x;

    #pragma unroll
    for(unsigned i = 0; i < NORM_THREAD_X; ++i)
    {
        uchar4 pixel = src[base + i * NORM_BLOCK_X];
        float I = (uint)pixel.x + pixel.y + pixel.z;
        dst[base + i * NORM_BLOCK_X] = (I + shift) * scale;
    }
}

void normalizeImage(float* dst, const uchar4* src, const uint2& dims)
{
    const dim3 blocks(dims.x * dims.y / NORM_PIXELS_PER_BLOCK);
    const dim3 threads(NORM_BLOCK_X);

    void* initialRange;
    CUDA_ERROR_CHECK(cudaGetSymbolAddress(&initialRange, INITIAL_RANGE));
    CUDA_ERROR_CHECK(cudaMemcpyToSymbol(normalRange, initialRange, sizeof(uint2)));

    normalizeImage1<<<blocks, threads>>>(src);
    normalizeImage2<<<blocks, threads>>>(dst, src);
}

#define DOWN_BLOCK_X 16
#define DOWN_BLOCK_Y 2

__global__ void downsample(float* out, uint2 dims)
{
    uint u = blockIdx.x * DOWN_BLOCK_X + threadIdx.x;
    uint v = blockIdx.y * DOWN_BLOCK_Y + threadIdx.y;

    float x = u / (float)dims.x;
    float y = v / (float)dims.y;

    out[dims.x * v + u] = tex2D(downSampler, x, y);
}

void downsample(float* dest, const uint2& destDims, const float* src, const uint2& srcDims)
{
    downSampler.normalized = true;
    downSampler.filterMode = cudaFilterModeLinear;

    cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float>();
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, downSampler, src, channelDesc, 
        srcDims.x, srcDims.y, srcDims.x * sizeof(float)));

    dim3 blocks(destDims.x / DOWN_BLOCK_X, destDims.y / DOWN_BLOCK_Y);
    dim3 threads(DOWN_BLOCK_X, DOWN_BLOCK_Y);

    downsample<<<blocks, threads>>>(dest, destDims);
}

#define WINDOW_RADIUS 3
#define WINDOW_WIDTH (WINDOW_RADIUS * 2 + 1)

#define EIGEN_BLOCK_SIZE 8

template<int i>
__device__ void gradMatUnrollX(float& gxx, float& gxy, float& gyy,
                               const int x, const int y, const int j)
{
    float gx = tex2D(gradxSampler, x + i, y + j);
    float gy = tex2D(gradySampler, x + i, y + j);
    gxx += gx * gx;
    gxy += gx * gy;
    gyy += gy * gy;
    gradMatUnrollX<i + 1>(gxx, gxy, gyy, x, y, j);
}

template<>
__device__ void gradMatUnrollX<WINDOW_RADIUS + 1>(float& gxx, float& gxy, float& gyy,
                                                  const int x, const int y, const int j)
{
}

template<int j>
__device__ void gradMatUnrollY(float& gxx, float& gxy, float& gyy,
                               const int x, const int y)
{
    gradMatUnrollX<-WINDOW_RADIUS>(gxx, gxy, gyy, x, y, j);
    gradMatUnrollY<j + 1>(gxx, gxy, gyy, x, y);
}

template<>
__device__ void gradMatUnrollY<WINDOW_RADIUS + 1>(float& gxx, float& gxy, float& gyy,
                                                  const int x, const int y)
{
}

__global__ void minEigenvalue
(
    float* eigenvalues,
    const uint width
)
{
    const int x = blockIdx.x * EIGEN_BLOCK_SIZE + threadIdx.x;
    const int y = blockIdx.y * EIGEN_BLOCK_SIZE + threadIdx.y;

    float gxx = 0.0f;
    float gxy = 0.0f;
    float gyy = 0.0f;

    // compute the 2x2 gradient matrix 
    //
    // [ gxx gxy ] 
    // [ gxy gyy ]
    //
    // from the sum of gradients in the surrounding window
    gradMatUnrollY<-WINDOW_RADIUS>(gxx, gxy, gyy, x, y);

    // compute the mininum eigenvalue of the gradient matrix
    float value = (gxx + gyy - sqrt((gxx - gyy) * (gxx - gyy) + 4 * gxy * gxy)) * 0.5f;
    eigenvalues[y * width + x] = value;
}

void minEigenvalue(float* dest, const float* gradx, const float* grady, const uint2& dims)
{
    cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float>();
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, gradxSampler, gradx, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, gradySampler, grady, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));

    const dim3 blocks(dims.x / EIGEN_BLOCK_SIZE, dims.y / EIGEN_BLOCK_SIZE);
    const dim3 threads(EIGEN_BLOCK_SIZE, EIGEN_BLOCK_SIZE);

    minEigenvalue<<<blocks, threads>>>(dest, dims.x);
}

__device__ float2 operator *(const float2& a, const float s) 
{
    return make_float2(a.x * s, a.y * s);
}

__device__ float2 operator +(const float2& a, const float2& b) 
{
    return make_float2(a.x + b.x, a.y + b.y);
}

#define ISCALE 255.0f
#define ISCALE2 (ISCALE * ISCALE)

#define STEP_FACTOR (-1.0f)
#define ITERATIONS 10
#define MIN_DETERMINANT 0.01f
#define MIN_DISPLACEMENT 0.1f
#define MAX_RESIDUE (10.0f * (WINDOW_WIDTH * WINDOW_WIDTH))

__global__ void trackFeatures
(
    FeaturePair* pairs,
    const float subsample
)
{
    float2 p1 = pairs[blockIdx.x].oldPos * subsample;
    float2 p2 = pairs[blockIdx.x].newPos * subsample;
    uint iter;

    if(p2.x < 0.0f)
    {
        return;
    }

    for(iter = 0; iter < ITERATIONS; ++iter)
    {
        float gxx = 0.0f;
        float gxy = 0.0f;
        float gyy = 0.0f;
        float2 error = make_float2(0.0f, 0.0f);

        for(int j = -WINDOW_RADIUS; j <= WINDOW_RADIUS; ++j)
        {
            for(int i = -WINDOW_RADIUS; i <= WINDOW_RADIUS; ++i)
            {
                // intensity difference
                float I1 = tex2D(image1Sampler, p1.x + i, p1.y + j);
                float I2 = tex2D(image2Sampler, p2.x + i, p2.y + j);
                float diff = I1 - I2;

                // gradient sum
                float gx1 = tex2D(gradx1Sampler, p1.x + i, p1.y + j);
                float gx2 = tex2D(gradx2Sampler, p2.x + i, p2.y + j);
                float gx = gx1 + gx2;
                float gy1 = tex2D(grady1Sampler, p1.x + i, p1.y + j);
                float gy2 = tex2D(grady2Sampler, p2.x + i, p2.y + j);
                float gy = gy1 + gy2;

                // 2x2 gradient matrix
                gxx += gx * gx;
                gxy += gx * gy;
                gyy += gy * gy;

                // 2x1 error vector
                error = error + make_float2(gx, gy) * diff;
            }
        }

        // scale to a reasonable factor
        error = error * (STEP_FACTOR * ISCALE2);
        gxx *= ISCALE2;
        gxy *= ISCALE2;
        gyy *= ISCALE2;

        // solve matrix equation
        float det = gxx * gyy - gxy * gxy;
        if(fabs(det) < MIN_DETERMINANT)
        {
            p2 = make_float2(-2.0f, -2.0f);
            break;
        }
        float invDet = 1.0f / det;
        float2 disp = make_float2(gyy * error.x - gxy * error.y, gxx * error.y - gxy * error.x) * invDet;

        p2 = p2 + disp;

        if(fabs(disp.x) < MIN_DISPLACEMENT || fabs(disp.y) < MIN_DISPLACEMENT)
        {
            break;
        }
    }

    if(iter < ITERATIONS)
    {
        // compute final residue
        float residue = 0.0f;
        for(int j = -WINDOW_RADIUS; j <= WINDOW_RADIUS; ++j)
        {
            for(int i = -WINDOW_RADIUS; i <= WINDOW_RADIUS; ++i)
            {
                // intensity difference
                float I1 = tex2D(image1Sampler, p1.x + i, p1.y + j);
                float I2 = tex2D(image2Sampler, p2.x + i, p2.y + j);
                residue += fabs(I1 - I2);
            }
        }
        if((residue * ISCALE) > MAX_RESIDUE)
        {
            p2 = make_float2(-3.0f, -3.0f);
        }
    }
    else
    {
        p2 = make_float2(-4.0f, -4.0f);
    }

    pairs[blockIdx.x].oldPos = p1;
    pairs[blockIdx.x].newPos = p2;
}

void trackFeatures(FeaturePair* pairs, const uint count,
                   const float* image1, const float* image2, 
                   const float* gradx1, const float* gradx2,
                   const float* grady1, const float* grady2,
                   const uint2& dims,
                   const float subsample)
{
    image1Sampler.filterMode = cudaFilterModeLinear;
    image2Sampler.filterMode = cudaFilterModeLinear;
    gradx1Sampler.filterMode = cudaFilterModeLinear;
    gradx2Sampler.filterMode = cudaFilterModeLinear;
    grady1Sampler.filterMode = cudaFilterModeLinear;
    grady2Sampler.filterMode = cudaFilterModeLinear;

    cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float>();
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, image1Sampler, image1, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, image2Sampler, image2, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, gradx1Sampler, gradx1, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, gradx2Sampler, gradx2, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, grady1Sampler, grady1, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));
    CUDA_ERROR_CHECK(cudaBindTexture2D(NULL, grady2Sampler, grady2, channelDesc, 
        dims.x, dims.y, dims.x * sizeof(float)));

    const dim3 blocks(count);
    const dim3 threads(1);

    trackFeatures<<<blocks, threads>>>(pairs, subsample);
}

NAMESPACE_END(gtd)